# jsb/databse/alchemy.py
#
#

""" jsb alchemy interface. """

## jsb imports

from jsb.lib.datadir import getdatadir
from jsb.lib.config import getmainconfig
from jsb.utils.locking import lockdec
from jsb.utils.exception import handle_exception
from jsb.imports import getsqlalchemy

## sqlalchemy imports

sqlalchemy = getsqlalchemy()

from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.associationproxy import association_proxy
from sqlalchemy.ext.orderinglist import ordering_list
from sqlalchemy import Text, Integer, Sequence, ForeignKey, DateTime
from sqlalchemy import create_engine, Column, String, Table
from sqlalchemy.orm import  scoped_session, sessionmaker, relation, eagerload
from sqlalchemy.orm import create_session as cs

## basic imports

import sqlalchemy, thread, os, time, logging

## debug settings

if getmainconfig()['debug']:
    logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

## locks

alchemylock = thread.allocate_lock()
dblocked = lockdec(alchemylock)
sessionlock = thread.allocate_lock()
sessionlocked = lockdec(sessionlock)
querylock = thread.allocate_lock()
querylocked = lockdec(querylock)
createlock = thread.allocate_lock()
createlocked = lockdec(createlock)

## defines

created = []

## create_all function

@createlocked
def create_all(plugname='all', base=None):
    """ create metadata on a plugin (or all). """
    logging.info('running create_all (%s)' % plugname)
    if plugname not in created:
        created.append(plugname)
        if not base: base = Base
        try: base.metadata.create_all()
        except Exception, ex: logging.warn('problem creating %s tables: %s' % (str(base), str(ex)))
    else: logging.warn('%s tables already created' % plugname)

## geturi function

def geturi(ddir=None, mainconfig=None):
    """  determine database URI from config file """
    d = ddir or getdatadir()
    config = mainconfig or getmainconfig()
    if not config['dburi']:
        if not 'sqlite' in config['dbtype'] and not 'mysql' in config['dbtype']:
            dburi = "%s://%s:%s@%s/%s" % (config['dbtype'], config['dbuser'], config['dbpasswd'], config['dbhost'], config['dbname'])
        elif 'mysql' in config['dbtype']:
            dburi = "%s://%s:%s@%s/%s?charset=utf8&use_unicode=0" % (config['dbtype'], config['dbuser'], config['dbpasswd'], config['dbhost'], config['dbname'])
        else:
            if not os.path.isdir(d + os.sep + 'db'): os.mkdir(d + os.sep + 'db')
            dburi = "sqlite:///%s/db/%s" % (ddir or getdatadir(), config['dbname'])
    else:
        dburi = config['dburi']
        try: dbtype = dburi.split(':')[0]
        except: logging.warn("can't extract db data from dburi") ; dbtype = 'unknown'
        if config['dbtype'] != dbtype: config['dbtype'] = dbtype ; config.save()
    logging.warn("database type is %s" % config['dbtype'])
    return dburi

## dbstart function

def dbstart(ddir=None, mainconfig=None, base=None):
    """ start the database connection setting Session and engine. """
    dburi = geturi(ddir, mainconfig)
    try: engine = create_engine(dburi, strategy='threadlocal', pool_recycle=60, pool_size=50, max_overflow=0)
    except: engine = create_engine(dburi, strategy='threadlocal', pool_recycle=3600)
    if not base: base = Base
    base.metadata.bind = engine
    create_all()
    Session = scoped_session(sessionmaker())
    Session.configure(bind=engine)
    if '///' in dburi: logging.warn('started database %s' % dburi)
    else: logging.warn('started database')
    return (Session, engine)

## vars

Base = declarative_base()
Session = engine = None


## startmaindb function

def startmaindb(ddir=None, mainconfig=None):
    """ start the main database. """
    global Session
    global engine

    Session, engine = dbstart(ddir, mainconfig)

## trans decorator

def trans(func, ismethod=True):
    """ transaction function attribute. """
    def dotransaction(*args, **kwargs):
        res = None
        arglist2 = list(args)  
        session = create_session()
        if not ismethod: arglist2.insert(0, session)
        else: arglist2.insert(1, session)
        session.begin(subtransactions=True)
        res = func(*arglist2, **kwargs)
        session.commit()
        session.flush()
        session.close()
        return res 

    @dblocked
    def transaction(*args, **kwargs):
        """ the tranasction wrapper .. works on methods. """
        try: res = dotransaction(*args, **kwargs) ; return res
        except sqlalchemy.exc.OperationalError, ex:   
            logging.warn('rolling back the session')
            Session.rollback()
            if 'server has gone away' in str(ex): logging.warn('mysql server has gone away') ; startmaindb()
        except Exception, ex:
            if "rollback" in str(ex): logging.warn('rolling back the session') ; Session.rollback()
            else: raise
        logging.warn("transaction failed: %s - retyring .." % str(ex))
        try:
            res = dotransaction(*args, **kwargs)
            try: Session.commit()
            except Exception, ex: Session.rollback() ; raise
            return res
        except Exception, ex:
            logging.warn('rolling back the session - giving up - %s failed due to %s' % (str(args), str(ex)))
            Session.rollback()
            handle_exception()
 
    return transaction

def transfunc(func):
    """ transaction wrapper for functions. """
    return trans(func, ismethod=False)

def rollback(session):
    """ rollback on provided session. """
    logging.warn('rollback on %s' % str(session))
    session.rollback()

def close(session):
    """ close provided session. """
    session.close()
    logging.warn('%s closed' % str(session))

@sessionlocked
def create_session():
    """ create a session ready for use. """
    session = Session()
    return session

#@querylocked
def query(q, session=None):
    """ do a query on the database. """
    session = create_session()
    res = session.query(q)
    session.close()
    return res

